#!/usr/bin/env python3
"""
基础提示生成器

定义所有提示生成器的基类接口
"""

import torch
import torch.nn as nn
import numpy as np
from typing import Dict, List, Tuple, Optional, Any
from abc import ABC, abstractmethod

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from utils.logger import get_logger


class BasePromptGenerator(nn.Module, ABC):
    """基础提示生成器抽象类"""
    
    def __init__(self, 
                 input_dim: int,
                 prompt_dim: int,
                 device: str = "cuda"):
        """
        初始化基础提示生成器
        
        Args:
            input_dim: 输入特征维度
            prompt_dim: 提示特征维度
            device: 设备
        """
        super().__init__()
        
        self.logger = get_logger(self.__class__.__name__)
        self.input_dim = input_dim
        self.prompt_dim = prompt_dim
        self.device = device
        
        # 基础投影层
        self.input_projection = nn.Linear(input_dim, prompt_dim)
        self.activation = nn.ReLU()
        self.dropout = nn.Dropout(0.1)
        
    @abstractmethod
    def forward(self, 
                image_features: torch.Tensor,
                context: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        前向传播
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            context: 上下文信息
            
        Returns:
            提示字典
        """
        pass
    
    @abstractmethod
    def generate_prompts(self, 
                        image: torch.Tensor,
                        metadata: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        生成提示
        
        Args:
            image: 输入图像 (B, C, H, W)
            metadata: 元数据
            
        Returns:
            提示字典
        """
        pass
    
    def extract_features(self, image: torch.Tensor) -> torch.Tensor:
        """
        提取图像特征
        
        Args:
            image: 输入图像 (B, C, H, W)
            
        Returns:
            图像特征 (B, C, H, W)
        """
        # 这里可以集成预训练的视觉编码器
        # 目前使用简单的卷积特征提取
        if not hasattr(self, 'feature_extractor'):
            self.feature_extractor = nn.Sequential(
                nn.Conv2d(3, 64, 3, padding=1),
                nn.ReLU(),
                nn.Conv2d(64, 128, 3, padding=1),
                nn.ReLU(),
                nn.AdaptiveAvgPool2d((32, 32))
            ).to(self.device)
        
        return self.feature_extractor(image)
    
    def compute_loss(self, 
                    predictions: Dict[str, torch.Tensor],
                    targets: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        计算损失
        
        Args:
            predictions: 预测结果
            targets: 目标结果
            
        Returns:
            损失值
        """
        total_loss = 0.0
        
        for key in predictions:
            if key in targets:
                if key == 'points':
                    # 点提示损失
                    pred_points = predictions[key]
                    target_points = targets[key]
                    loss = nn.MSELoss()(pred_points, target_points)
                elif key == 'boxes':
                    # 框提示损失
                    pred_boxes = predictions[key]
                    target_boxes = targets[key]
                    loss = nn.MSELoss()(pred_boxes, target_boxes)
                elif key == 'masks':
                    # 掩码提示损失
                    pred_masks = predictions[key]
                    target_masks = targets[key]
                    loss = nn.BCEWithLogitsLoss()(pred_masks, target_masks)
                else:
                    # 默认L2损失
                    loss = nn.MSELoss()(predictions[key], targets[key])
                
                total_loss += loss
        
        return total_loss
    
    def save_checkpoint(self, save_path: str):
        """保存检查点"""
        checkpoint = {
            'model_state_dict': self.state_dict(),
            'input_dim': self.input_dim,
            'prompt_dim': self.prompt_dim,
            'device': self.device
        }
        torch.save(checkpoint, save_path)
        self.logger.info(f"模型检查点已保存: {save_path}")
    
    def load_checkpoint(self, load_path: str):
        """加载检查点"""
        checkpoint = torch.load(load_path, map_location=self.device)
        self.load_state_dict(checkpoint['model_state_dict'])
        self.logger.info(f"模型检查点已加载: {load_path}")
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息"""
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
        
        return {
            'model_name': self.__class__.__name__,
            'input_dim': self.input_dim,
            'prompt_dim': self.prompt_dim,
            'total_params': total_params,
            'trainable_params': trainable_params,
            'device': self.device
        }
    
    def enable_self_prompt(self, enable: bool = True):
        """启用/禁用自提示机制"""
        if hasattr(self, 'self_prompt_enabled'):
            self.self_prompt_enabled = enable
            self.logger.info(f"自提示机制 {'启用' if enable else '禁用'}")
        else:
            self.logger.warning("当前模型不支持自提示机制")
    
    def get_self_prompt_status(self) -> bool:
        """获取自提示状态"""
        return getattr(self, 'self_prompt_enabled', False)
